#ifndef ELEMENTS_ELEMENT_CONTAINER_HPP
#define ELEMENTS_ELEMENT_CONTAINER_HPP

#include <memory>
#include <mutex>
#include <shared_mutex>
#include <unordered_map>
#include <vector>

namespace Elements
{

template <typename T>
class ElementContainer final
{
public:
    std::vector<unsigned> ids() const
    {
        std::shared_lock locker(m_lock);
        std::vector<unsigned> result;
        result.reserve(m_elements.size());
        for (auto [id, element] : m_elements)
            result.push_back(id);
        return result;
    }

    std::vector<std::shared_ptr<T>> elements() const
    {
        std::shared_lock locker(m_lock);
        std::vector<std::shared_ptr<T>> result;
        result.reserve(m_elements.size());
        for (auto [id, element] : m_elements)
            result.push_back(element);
        return result;
    }

    std::shared_ptr<T> cloneElement(unsigned id) const
    {
        std::shared_lock locker(m_lock);
        auto elementIterator = m_elements.find(id);
        if (elementIterator == m_elements.end())
            return nullptr;

        return std::shared_ptr<T>(
            static_cast<T*>(elementIterator->second->clone())
        );
    }

    std::shared_ptr<T> element(unsigned id) const
    {
        std::unique_lock locker(m_lock);
        auto elementIterator = m_elements.find(id);
        if (elementIterator == m_elements.end())
            return nullptr;

        return elementIterator->second;
    }

    template<typename V>
    std::vector<std::shared_ptr<V>> elementsCast() const
    {
        std::shared_lock locker(m_lock);
        std::vector<std::shared_ptr<V>> result;
        result.reserve(m_elements.size());
        for (auto [id, element] : m_elements)
            result.push_back(std::static_pointer_cast<V>(element));
        return result;
    }

    template<typename V, typename P>
    std::vector<std::shared_ptr<V>> elementsCastByPredicate(P predicate) const
    {
        std::unique_lock locker(m_lock);
        std::vector<std::shared_ptr<V>> result;
        auto current = m_elements.cbegin();
        const auto end = m_elements.cend();
        while (end != current)
        {
            if (predicate(current->second))
                result.push_back(std::static_pointer_cast<V>(current->second));
            ++current;
        }
        return result;
    }

    template<typename V>
    std::shared_ptr<V> elementCast(unsigned id) const
    {
        std::unique_lock locker(m_lock);
        auto elementIterator = m_elements.find(id);
        if (elementIterator == m_elements.end())
            return nullptr;

        return std::static_pointer_cast<V>(elementIterator->second);
    }

    template<typename P>
    std::vector<std::shared_ptr<T>> elementsByPredicate(P predicate) const
    {
        std::unique_lock locker(m_lock);
        std::vector<std::shared_ptr<T>> result;
        auto current = m_elements.cbegin();
        const auto end = m_elements.cend();
        while (end != current)
        {
            if (predicate(current->second))
                result.push_back(current->second);
            ++current;
        }
        return result;
    }

    template<typename P>
    std::shared_ptr<T> elementByPredicate(P predicate) const
    {
        std::unique_lock locker(m_lock);
        auto current = m_elements.cbegin();
        const auto end = m_elements.cend();
        while (end != current)
        {
            if (predicate(current->second))
                return current->second;
            ++current;
        }
        return nullptr;
    }

    void add(std::shared_ptr<T> element)
    {
        if (nullptr == element)
            return;

        std::unique_lock locker(m_lock);
        m_elements[element->id()] = element;
    }

    template<typename V>
    void insertCast(const std::vector<std::shared_ptr<V>>& elements)
    {
        if (elements.empty())
            return;

        std::unique_lock locker(m_lock);
        for (auto element : elements)
            m_elements[element->id()] = std::static_pointer_cast<T>(element);
    }

    void remove(unsigned id)
    {
        std::unique_lock locker(m_lock);
        m_elements.erase(id);
    }

    template<typename Container>
    void remove(const Container& ids)
    {
        std::unique_lock locker(m_lock);
        for (const auto id : ids)
            m_elements.erase(id);
    }

    unsigned size() const
    {
        std::shared_lock locker(m_lock);
        return m_elements.size();
    }

    void clear()
    {
        std::unique_lock locker(m_lock);
        return m_elements.clear();
    }

private:
    mutable std::shared_mutex m_lock;
    std::unordered_map<unsigned, std::shared_ptr<T>> m_elements;
};

} // namespace Elements

#endif // ELEMENTS_ELEMENT_CONTAINER_HPP
